import numpy as np
import tvm
import tvm.relay as relay
from tvm.contrib import graph_runtime
import mxnet as mx
from mxnet import ndarray as nd

class TVMGetFeatureModel():
    def __init__(self,):
        ctx = tvm.gpu(0)
        # load the module back.
        loaded_json = open("model/tvm/deploy_graph.json").read()
        loaded_lib = tvm.runtime.load_module("model/tvm/deploy_lib.so")
        loaded_params = bytearray(open("model/tvm/deploy_param.params", "rb").read())

        module = graph_runtime.create(loaded_json, loaded_lib, ctx)
        module.load_params(loaded_params)
        self.module=module

    def get_feature(self, img):
        #try:
        print(img.shape)
        input_data = np.array(img)
        input_data = tvm.nd.array(input_data.astype("float32"))
        self.module.run(data=input_data)
        feature = self.module.get_output(0).asnumpy()[0]
        print(feature.shape)
        return None, feature
        #except Exception as e1:
            #print('-------get feature error')
            #return (90099, '服务超时'), None